Python yolov5加Flask实现发布图像识别服务

yolov5加Flask实现发布图像识别服务

使用Flask作为后端，简单方便，接收到前端传过来的base64,再转为图片，输入到YOLOv5模型当中进行识别，然后会得到识别完成的base64图像，再返回给前端。
使用yolov5s.pt模型。直接使用torch.hub.load()进行模型加载。从本地加载可以避免由于github下载导致的问题。图像传递的过程为web->flask。传递方式都为base64。

加载模型，输入图片进行detect，返回图片的base64编码。

from flask import Flask, jsonify, request, abort
import json
import base64
import numpy as np

import torch
import os
from PIL import Image
import cv2
from io import BytesIO, StringIO


app = Flask(__name__)

@app.route('/ana', methods=['GET', 'POST'])
def call_analysis():
    print("Hello, World!")
 
    data1 = request.data   #----获取的是字符串
    #print(data1)
    data2 = request.get_data()  #----获取的是字符串
    #print(data2)
    j_data =  json.loads(data2) #-----load将字符串解析成json
    # print(j_data)
    # print(j_data['base64']) 
    
    # return jsonify(j_data)
    return getbase64(j_data['base64'])

def getbase64(imgbase64):
    # print(img) 
    # os.environ["CUDA_VISIBLE_DEVICES"]
    # model = torch.hub.load('./', 'custom', path='./yolov5s.pt', source='local')
    model = torch.hub.load('./yolov5-master/', 'custom', './yolov5-master/runs/train/exp2/weights/last.pt', source='local')
    # model.cuda()
    # print(model)
    # img1 = Image.open(BytesIO(img))

    
    # 传入为RGB格式下的base64，传出为RGB格式的numpy矩阵
    byte_data = base64.b64decode(imgbase64)#将base64转换为二进制
    encode_image = np.asarray(bytearray(byte_data), dtype="uint8")# 二进制转换为一维数组
    img_array = cv2.imdecode(encode_image, cv2.IMREAD_COLOR)# 用cv2解码为三通道矩阵
    img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)# BGR2RGB
 
    img1 = img_array
    imgs = [img1]
    #img2 = cv2.imread('./data/images/zidane.jpg')[:,:,::-1]
    #imgs = [img2]

    results = model(imgs, size=640)
    print('识别成功')
    results.ims
    results.render()
    base64_image = ''
    for img in results.ims:
        buffered = BytesIO()
        img_base64 = Image.fromarray(img)
        img_base64.save(buffered, format='JPEG')
        base64_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
        # with open('./result.txt', 'a+') as f:
        #     f.write(base64_image)
        #     f.close()
    print(results)
    results.save()
    base64_image = 'data:image/jpeg;base64,%s' % base64_image
    return base64_image
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug = False)

注意：torch.hub.load() 方法中的路径问题

一辈子很短，努力的做好两件事就好；
第一件事是热爱生活，好好的去爱身边的人；
第二件事是努力学习，在工作中取得不一样的成绩，实现自己的价值，而不是仅仅为了赚钱；